import sys
from PyQt5.QtGui import QIcon, QFont
from PyQt5.QtWidgets import QApplication, QWidget, \
    QLabel, QLineEdit, QPushButton, QFileDialog, QHBoxLayout, QVBoxLayout, QComboBox, QProgressBar, QMessageBox
from Progress import MainThread
from PyQt5.QtCore import Qt

class MainWindow(QWidget):

    def __init__(self):
        super().__init__()

        self.initUI()

    def initUI(self):

        with open('style.qss', 'r') as f:
            stylesheet = f.read()
            self.setStyleSheet(stylesheet)

        self.setGeometry(500, 300, 700, 350)  # 设置窗口大小
        self.setWindowTitle('Video-SR')  # 设置窗口名称
        self.setWindowIcon(QIcon('icon.png'))  # 设置窗口图片

        # Create labels and text boxes for video input and output paths
        self.video_input_label = QLabel(self)
        self.video_input_label.setText('Video Input Path:')
        self.video_input_label.setFont(QFont("Arial", 12))
        self.video_input_textbox = QLineEdit(self)

        self.video_output_label = QLabel(self)
        self.video_output_label.setText('Video Output Dir:')
        self.video_output_label.setFont(QFont("Arial", 12))
        self.video_output_textbox = QLineEdit(self)

        # Create buttons to select video input and output paths
        self.video_input_button = QPushButton('Browse', self)
        self.video_output_button = QPushButton('Browse', self)

        # Connect button clicks to methods for selecting video input and output paths
        self.video_input_button.clicked.connect(self.select_video_input_path)
        self.video_output_button.clicked.connect(self.select_video_output_path)

        # Create horizontal boxes for each set of label, textbox, and button
        hbox_input = QHBoxLayout()
        hbox_input.addWidget(self.video_input_label)
        hbox_input.addWidget(self.video_input_textbox)
        hbox_input.addWidget(self.video_input_button)

        hbox_output = QHBoxLayout()
        hbox_output.addWidget(self.video_output_label)
        hbox_output.addWidget(self.video_output_textbox)
        hbox_output.addWidget(self.video_output_button)

        # Create a vertical box to hold the two horizontal boxes
        vbox = QVBoxLayout()
        vbox.addLayout(hbox_input)
        vbox.addSpacing(10)  # 添加垂直间距
        vbox.addLayout(hbox_output)
        vbox.addSpacing(10)

        # Add a QComboBox control for selecting whether or not to use GPU
        self.gpu_label = QLabel(self)
        self.gpu_label.setText('Use GPU:')
        self.gpu_label.setFont(QFont("Arial", 12))  # 设置字体和字号
        self.gpu_combo_box = QComboBox(self)
        self.gpu_combo_box.addItem('Yes')
        self.gpu_combo_box.addItem('No')

        gpu_hbox = QHBoxLayout()
        gpu_hbox.addWidget(self.gpu_label)
        gpu_hbox.addWidget(self.gpu_combo_box)

        vbox.addLayout(gpu_hbox)

        # Add a QComboBox control for selecting the model
        self.model_label = QLabel(self)
        self.model_label.setText('Model:')
        self.model_label.setFont(QFont("Arial", 12))
        self.model_combo_box = QComboBox(self)
        self.model_combo_box.addItem('PP-MSVSR_reds_x4')
        self.model_combo_box.addItem('EDVR_M_wo_tsa_SRx4')
        self.model_combo_box.addItem('BasicVSR_reds_x4')

        model_hbox = QHBoxLayout()
        model_hbox.addWidget(self.model_label)
        model_hbox.addWidget(self.model_combo_box)

        vbox.addLayout(model_hbox)

        # Add progress bar and label
        self.progress_label = QLabel("Infer: ", self)
        self.progress_label.setFont(QFont("Arial", 12))
        self.progress_bar = QProgressBar(self)
        self.progress_bar.setRange(0, 100)

        progress_hbox = QHBoxLayout()
        progress_hbox.addWidget(self.progress_label)
        progress_hbox.addWidget(self.progress_bar)
        vbox.addLayout(progress_hbox)


        # Add start and stop buttons
        self.start_button = QPushButton("Start", self)
        self.start_button.setFont(QFont("Arial", 14))
        self.start_button.clicked.connect(self.start_thread)
        vbox.addWidget(self.start_button, alignment=Qt.AlignHCenter)  # 将 start 按钮添加到 vbox 中，并水平居中对齐

        self.setLayout(vbox)
        self.show()

    def start_thread(self):

        input_path = self.video_input_textbox.text().strip()
        output_path = self.video_output_textbox.text().strip()
        model_name = self.model_combo_box.currentText()
        # 路径是否为空进行检查
        if input_path == "":
            QMessageBox.warning(self, "Warning", "Please select a video input path")
            return
        if output_path == "":
            QMessageBox.warning(self, "Warning", "Please select a video output path")
            return

        if self.gpu_combo_box.currentText() == "Yes":
            use_gpu = True
        else:
            use_gpu = False

        self.thread = MainThread(input_path, output_path, model_name, use_gpu)
        self.thread.signal.connect(self.updateProgressBar)
        self.thread.start()

    def updateProgressBar(self, msg):
        if msg == "finished":
            self.progress_bar.setValue(100)
            QMessageBox.information(self, "Information", "Super Resolution Completed!")
            self.progress_bar.reset()
        else:
            try:
                progress = int(msg)
                self.progress_bar.setValue(progress)
            except ValueError:
                QMessageBox.warning(self, "Warning", msg)

    def select_video_input_path(self):
        """
        Open a dialog box to select the video input path
        """
        file_dialog = QFileDialog()
        file_dialog.setNameFilter("Videos (*.mp4)")
        file_dialog.selectNameFilter("Videos (*.mp4)")

        if file_dialog.exec_():
            filenames = file_dialog.selectedFiles()
            self.video_input_textbox.setText(filenames[0])

    def select_video_output_path(self):
        """
        Open a dialog box to select the video output path
        """
        file_dialog = QFileDialog()
        file_dialog.setFileMode(QFileDialog.DirectoryOnly)

        if file_dialog.exec_():
            directory = file_dialog.selectedFiles()[0]
            self.video_output_textbox.setText(directory)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = MainWindow()
    sys.exit(app.exec_())
